Move existing PPR sampler work off event loop - #708
Merged
Conversation
mkolodner-sc
requested review from
kmontemayor2-sc,
nshah-sc,
svij-sc,
xgao4-sc,
yliu2-sc and
zfan3-sc
as code owners
July 15, 2026 18:54
This was referenced Jul 15, 2026
mkolodner-sc
force-pushed
the
mkolodner-sc/pr701-ppr-cleanup
branch
from
July 15, 2026 20:02
5197b5d to
7f6d9c0
Compare
mkolodner-sc
force-pushed
the
mkolodner-sc/pr701-ppr-event-loop
branch
from
July 15, 2026 20:02
fc410ab to
2b85456
Compare
mkolodner-sc
force-pushed
the
mkolodner-sc/pr701-ppr-cleanup
branch
from
July 15, 2026 21:24
7f6d9c0 to
19c0d92
Compare
mkolodner-sc
force-pushed
the
mkolodner-sc/pr701-ppr-event-loop
branch
from
July 15, 2026 21:24
2b85456 to
ecb394c
Compare
mkolodner-sc
force-pushed
the
mkolodner-sc/pr701-ppr-cleanup
branch
from
July 15, 2026 21:40
19c0d92 to
f7aec06
Compare
mkolodner-sc
force-pushed
the
mkolodner-sc/pr701-ppr-event-loop
branch
from
July 15, 2026 21:40
ecb394c to
0047b2f
Compare
mkolodner-sc
force-pushed
the
mkolodner-sc/pr701-ppr-event-loop
branch
3 times, most recently
from
July 15, 2026 23:04
60a0cea to
ff2d3ea
Compare
kmontemayor2-sc
approved these changes
Jul 15, 2026
kmontemayor2-sc
left a comment
kmontemayor2-sc
left a comment
Collaborator
There was a problem hiding this comment.
LGTM! Nice work Matt :)
Do you think there's similar low-hanging (but hard to find!) fruit here for the khop sampler?
E.g. is it possible to put this bit in the background? https://github.com/alibaba/graphlearn-for-pytorch/blob/main/graphlearn_torch/python/distributed/dist_neighbor_sampler.py#L297
Or do we need to also write a CPP version of that _sample_from_nodes so we can force-detach the gil? May not be difficult? WDYT?
mkolodner-sc
force-pushed
the
mkolodner-sc/pr701-ppr-event-loop
branch
from
July 16, 2026 18:46
ff2d3ea to
d0f9cea
Compare
mkolodner-sc
force-pushed
the
mkolodner-sc/pr701-ppr-cleanup
branch
2 times, most recently
from
July 16, 2026 19:15
2dd7085 to
68d6d62
Compare
mkolodner-sc
force-pushed
the
mkolodner-sc/pr701-ppr-event-loop
branch
2 times, most recently
from
July 16, 2026 19:21
33c8dde to
1b0ea2a
Compare
mkolodner-sc
force-pushed
the
mkolodner-sc/pr701-ppr-cleanup
branch
2 times, most recently
from
July 16, 2026 20:01
3d34854 to
f5c430d
Compare
mkolodner-sc
force-pushed
the
mkolodner-sc/pr701-ppr-event-loop
branch
from
July 16, 2026 20:01
1b0ea2a to
ed7387c
Compare
yliu2-sc
approved these changes
Jul 22, 2026
mkolodner-sc
force-pushed
the
mkolodner-sc/pr701-ppr-event-loop
branch
from
July 27, 2026 10:20
ed7387c to
457c30b
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Moves the existing distributed PPR sampler's CPU-heavy C++ work off the asyncio event-loop path.
This PR is limited to the existing/non-typed PPR path. Typed-channel PPR enablement is left for the next PR in the stack.
Stack
Reasoning for Change
PPR sampling runs several CPU-heavy C++ steps while async graph-store/RPC work is also in flight. When those steps run on the event-loop thread or hold the GIL longer than necessary, they can delay RPC completion callbacks and reduce concurrency across sampler coroutines.
This change keeps the PPR math and returned sequences unchanged while moving construction, queue draining, residual pushes, and extraction through executor/GIL-release boundaries where possible.
Changes
PPRForwardPushconstruction, queue draining, residual push, and extraction in the pybind layer.drain_queue,push_residuals, and top-k extraction throughrun_in_executorfrom the Python sampler..to(device)before one-hop fetches; extracted tensors are still moved to the sampler device at the existing output boundary.